import psycopg2
from contextlib import contextmanager
from config import PGVECTOR_URL

class DatabaseManager:
    def __init__(self, db_url):
        self.db_url = db_url

    @contextmanager
    def get_cursor(self):
        conn = None
        try:
            # 创建数据库连接
            conn = psycopg2.connect(self.db_url)
            cursor = conn.cursor()
            yield cursor
            conn.commit()  # 提交事务
        except Exception as e:
            if conn:
                conn.rollback()  # 出现异常时回滚事务
            raise e
        finally:
            if cursor:
                cursor.close()
            if conn:
                conn.close()

    def execute_query(self, query, params=None):
        with self.get_cursor() as cur:
            cur.execute(query, params)
            return cur.fetchall()

    def insert(self, table, columns, values):
        placeholders = ', '.join(['%s'] * len(values))
        columns_str = ', '.join(columns)
        query = f"INSERT INTO {table} ({columns_str}) VALUES ({placeholders})"
        # print("executing sql: ", query, values)
        with self.get_cursor() as cur:
            cur.execute(query, values)

    def update(self, table, set_data, condition):
        set_clause = ', '.join([f"{key}=%s" for key in set_data])
        condition_clause = ' AND '.join([f"{key}=%s" for key in condition])
        query = f"UPDATE {table} SET {set_clause} WHERE {condition_clause}"
        # print("executing sql: ", query, tuple(set_data.values()) + tuple(condition.values()))
        with self.get_cursor() as cur:
            cur.execute(query, tuple(set_data.values()) + tuple(condition.values()))

    def delete(self, table, condition):
        condition_clause = ' AND '.join([f"{key}=%s" for key in condition])
        query = f"DELETE FROM {table} WHERE {condition_clause}"
        # print("executing sql: ", query, tuple(condition.values()))
        with self.get_cursor() as cur:
            cur.execute(query, tuple(condition.values()))

    def fetch_all(self, table, columns='*', condition=None):
        condition_clause = ' WHERE ' + ' AND '.join([f"{key}=%s" for key in condition]) if condition else ''
        query = f"SELECT {columns} FROM {table} {condition_clause}"
        # print(query, tuple(condition.values()) if condition else None)
        return self.execute_query(query, tuple(condition.values()) if condition else None)

database = "/postgres"
db_url = PGVECTOR_URL + database
db_manager = DatabaseManager(db_url)


# 使用示例
if __name__ == "__main__":

    # 接口使用示例 传入的参数如下
    # db_manager.insert('public.task_status', ['task_id', 'status_id', 'status'], ['100001', '0', 'running'])
    #
    # db_manager.update('public.task_status', {"status_id": "1"}, {"task_id": "100001"})
    # db_manager.update('public.task_status', {"status_id": "1", "status": "success"}, {"task_id": "100001"})
    #
    # db_manager.delete('public.task_status', {'task_id': '11111111'})
    #
    # res = db_manager.fetch_all('public.task_status')
    # print(res)
    #
    # res = db_manager.fetch_all('public.task_status', "task_id,status_id")
    # print(res)
    #
    # res = db_manager.fetch_all('public.task_status', condition={"task_id": "100001"})
    # print(res)

    res = db_manager.fetch_all('public.task_status', columns="task_id", condition={"status": "success"})
    print(res)
    for i in res:
        print(i[0])


